import numpy as np
from sklearn.neighbors import NearestNeighbors

def custom_metric(arr1, arr2):
    """保持张量计算连续性"""
    # 确保输入为PyTorch张量
    if not isinstance(arr1, torch.Tensor):
        arr1 = torch.from_numpy(arr1)
    if not isinstance(arr2, torch.Tensor):
        arr2 = torch.from_numpy(arr2)
    
    # 保持计算图完整性
    centroid_diff = calculate_centroid(arr1) - calculate_centroid(arr2)
    return torch.linalg.norm(centroid_diff)

def calculate_centroid(points):
    """兼容PyTorch张量和Numpy数组的质心计算"""
    if isinstance(points, torch.Tensor):
        return torch.mean(points, dim=0)
    return np.mean(points, axis=0)

def adjacency_edges(points, k=2):
    """生成邻接边坐标"""
    nbrs = NearestNeighbors(n_neighbors=k+1).fit(points)
    _, indices = nbrs.kneighbors(points)
    
    edges = []
    for i in range(len(points)):
        for j in indices[i][1:]:  # 跳过自身
            edges.append((*points[i], *points[j]))
    return edges

def geometry_distance(points1, points2):
    """支持张量输入并返回浮点数（梯度兼容版）"""
    # 确保输入为PyTorch张量
    if not isinstance(points1, torch.Tensor):
        points1 = torch.tensor(points1, dtype=torch.float32)
    if not isinstance(points2, torch.Tensor):
        points2 = torch.tensor(points2, dtype=torch.float32)
    
    # 新增点数校验
    if points1.shape[0] != points2.shape[0]:
        raise ValueError(f"点集数量不匹配：左图{points1.shape[0]}点 vs 右图{points2.shape[0]}点")
    
    res = geometry_dist(points1.unsqueeze(0), points2.unsqueeze(0), c='dummy')[0]
    return res  # 返回可求导的张量


#---------------------------------copied and adapted from ht2t--------------------------------
import torch
from torch.nn.functional import normalize
def seq_euclidean_dist(x, y, eps=1e-8):
    # return torch.sqrt(torch.sum((x - y) ** 2, dim=-1) + eps)
    
    #squared
    return torch.sum((x - y) ** 2, dim=-1)

def normalize_matrix(m, dims=(1,2), eps=1e-8):
    frob_norm = torch.linalg.norm(m, ord='fro', dim=dims, keepdim=True)
    normalized_m = m / (frob_norm + eps)
    return normalized_m

def morse_new(A, D):
    #(batch_size, colnum, colnum)
    T0 = A
    T1 = torch.bmm(A, D)
    T1_1 = torch.bmm(D, A)
    T2 = torch.bmm(T1, T1_1)
    T3 = torch.bmm(torch.bmm(T2, T1_1), T1)
    #for debug:
    # print('T0', T0);print('T1', T1);print('T1_1', T1_1);print('T2', T2);print('T3', T3)
    return torch.stack([T0,T1,T2,T3], dim=1)

def seq_hie_morse_new(tree_seq, c):
    set_size = tree_seq.shape[1]
    # A = seq_poincare_dist(tree_seq.unsqueeze(2).expand(-1, -1, set_size, -1), tree_seq.unsqueeze(1).expand(-1, set_size, -1, -1), c=c)
    A = seq_euclidean_dist(tree_seq.unsqueeze(2).expand(-1, -1, set_size, -1), tree_seq.unsqueeze(1).expand(-1, set_size, -1, -1))
    D = torch.diag_embed(A.sum(dim=2))
    A, D = normalize_matrix(A, dims=(1,2)), normalize_matrix(D, dims=(1,2))
    M = morse_new(A, D)
    return  M

def morse2vec(M_seq, w):
    #M_seq.shape (batch_size, 5, node_num, node_num)
    w=w.unsqueeze(-1).unsqueeze(-1).unsqueeze(0) #(1, 5, 1, 1)
    M_seq = torch.flatten(w * M_seq, start_dim=1)
    return M_seq

def gen_morse_vector(tree_seq, c, w=None, morse_mp=None):
    # w = w if w is not None else torch.tensor([0.25,0.25,0.25,0.25]).cuda()
    w = w if w is not None else torch.tensor([0.25,0.25,0.25,0.25])
    M = seq_hie_morse_new(tree_seq, c)
    # M = morse_mp(morse2vec(M, w))
    M = morse2vec(M, w)
    if morse_mp is not None:
        M = morse_mp(M)
    # M = normalize(M, p=2, dim=1)
    return M

def geometry_dist_matrix(tree_seq1, tree_seq2, c, w=None, morse_mp=None):
    M1, M2 = gen_morse_vector(tree_seq1, c, w=w, morse_mp=morse_mp), gen_morse_vector(tree_seq2, c, w=w, morse_mp=morse_mp)
    # dist_matrix = 2 - 2 * (M1 @ M2.T)
    dist_matrix = 2 - 2 * (M1 @ M2.T)
    return dist_matrix

def geometry_dist(tree_seq1, tree_seq2, c, w=None, morse_mp=None):
    M1, M2 = gen_morse_vector(tree_seq1, c, w=w, morse_mp=morse_mp), gen_morse_vector(tree_seq2, c, w=w, morse_mp=morse_mp)
    # dist = 2 - 2 * (M1 * M2).sum(dim=-1)
    dist = ((M1 - M2)**2).sum(dim=-1)
    return dist

if __name__ == '__main__':
    # 创建张量时明确设置 requires_grad=True
    a = torch.tensor([[-1.,1],[1,1], [1,-1], [-1, -1]], 
                    dtype=torch.float32, 
                    requires_grad=True)  # 确保是叶子张量
    b = torch.tensor([[-2.,2],[2,2], [2, -2], [-2, -2]], 
                    dtype=torch.float32, 
                    requires_grad=True)  # 确保是叶子张量
    
    loss = -custom_metric(a + torch.tensor([2, 2]), b)
    loss.backward()
    print("loss:", loss.item())  # 输出损失值
    print("a.is_leaf:", a.is_leaf)  # 输出: True
    print("b.is_leaf:", b.is_leaf)  # 输出: True
    print("a.grad:", a.grad)  # 正确梯度
    print("b.grad:", b.grad)  # 正确梯度